# -*- coding: utf-8 -*-
'''
Moduł zawierający implementację algorytmu Needlemana-Wunscha dla 3 sekwencji.

'''
# from __future__ import #print_function
import FastaParser
from PyQt4.QtCore import Qt, QThread, SIGNAL, QMutex
import os.path as path
import copy
from numpy import eye, inf, argmax, ndarray
from os import remove

from memory_profiler import profile
import timeit

## Klasa reprezentująca algorytm Needlemana-Wunscha z możliwością obliczenia dopasowania w osobnym wątku.
class SequenceAlignment(QThread):

    ##Kod wyboru macierzy podobieństwa IUB
    IUB = 0
    ##Kod wyboru macierzy podobieństwa CLUSTALW
    CLUSTALW = 1
    ##Kod wyboru własnej macierzy
    USER = 2
    ##Minimalny możliwy do wybrania numer macierzy
    MIN_MATRIX_NUM = IUB
    ##Maksymalny możliwy do wybrania numer macierzy
    MAX_MATRIX_NUM = USER
    ##Maksymalna możliwa do ustawienia wartość kary
    PENALTY_MAX = 0.0
    ##Minimalna wartość w macierzy prawdopodobieństwa
    MATRIX_MIN_VAL = -99.99
    ##Maksymalna wartość w macierzy prawdopodobieństwa
    MATRIX_MAX_VAL = 99.99
    
    ##Lista nukleotydow
    NUCLEOTIDE_LIST = ['A', 'G', 'C', 'T', '-']
    ##Lista przypisań kolorów do nukleotydów
    NUCLEOTIDE_COLORS = [Qt.red, Qt.yellow, Qt.blue, Qt.green, Qt.white]
    ##Wymiar macierzy podobieństwa
    DIM = len(NUCLEOTIDE_LIST)
    ##Definicja macierzy podobieństwa IUB
    IUB_MATRIX = 1.9*eye(DIM, DIM)
    ##Definicja macierzy podobieństwa CLUSTALW
    CLUSTALW_MATRIX = eye(DIM, DIM)

    ##Konstruktor klasy, dziedziczy po QThread
    #@param parent rodzic klasy
    def __init__(self, parent = None):
        QThread.__init__(self, parent)
        ##Numer wybranej macierz podobieństwa
        self.matrixNum = SequenceAlignment.CLUSTALW
        ##Ustalona wartość kary za przerwę
        self.penaltyVal = SequenceAlignment.PENALTY_MAX
        ##Macierz podobieństwa
        self.matrix = SequenceAlignment.CLUSTALW_MATRIX
        self._fillGapInMatrix()
        
        ##Pierwsza wczytana sekwencja
        self.seq1 = ''
        ##Druga wczytana sekwencja
        self.seq2 = ''
        ##Trzecia wczytana sekwencja
        self.seq3 = ''
        ##Nazwa pierwszej wczytanej sekwencji
        self.name1 = ''
        ##Nazwa drugiej wczytanej sekwencji
        self.name2 = ''
        ##Nazwa trzeciej wczytanej sekwencji
        self.name3 = ''
        ##Lista zawierająca ’*’, gdy wszystkie elementy sekwencji takie same lub ’ ’, gdy różne
        self.seqColumns = []
        ##Mutex na wartość zmiennej decydującej o zakończeniu przetwarzania
        self.stopProcessingMutex = QMutex()
        ##Wartość ustawiana przez gui, gdy należy przerwać przetwarzanie
        self.stopProcessingVal = 0
        
        
    ##Wczytanie sekwencji do porównania z pliku
    #@param filePath ścieżka do pliku
    #@return listę par [nazwa, sekwencja] oraz listę z danymi, czy wszystkie nukleotydy na danej pozycji takie same 
    def readSequences(self, filePath):
        try:
            with open(filePath, 'r') as fp:
                sequencesList = []

                for pair in FastaParser.read_fasta(fp):
                    sequencesList.append(pair)
                    
                #jezeli jest za duzo lub za malo sekwencji
                if len(sequencesList) != 3:
                    return None, None
                
                #jezeli podano pusta sekwencje lub zawierajaca znaki spoza NUCLEOTIDE_LIST
                for (name, seq) in sequencesList:
                    if name == "" or seq == "" or any(map((lambda x: x not in SequenceAlignment.NUCLEOTIDE_LIST), seq)):
                        return None, None
                ##Pierwsza wyrównana sekwencja
                self.ress1 = []
                ##Druga wyrównana sekwencja
                self.ress2 = []
                ##Trzecia wyrównana sekwencja
                self.ress3 = []
                ##Lista zawierająca '*', gdy wszystkie elementy wyrównanych sekwencji takie same, ' ' gdy różne
                self.ressColumns = []
                ##Wartość wyrównania sekwencji
                self.alignmentValue = -inf       
                self.name1 = sequencesList[0][0]
                self.name2 = sequencesList[1][0]
                self.name3 = sequencesList[2][0]      
                self.seq1 = sequencesList[0][1]
                self.seq2 = sequencesList[1][1]
                self.seq3 = sequencesList[2][1]
                self.findExactInput()
                return sequencesList, self.seqColumns
        except:
            return None, None
    
    ##Funkcja zapisująca wyrównane sekwencje do pliku
    #@param filePath ścieżka do pliku
    #@return: filePath, gdy poprawnie zapisane, None gdy błąd
    def saveAlignment(self, filePath):
        isOverwriting = path.isfile(filePath)
        try:
            with open(filePath, 'w') as fp:
                FastaParser.save_clustal(fp, [self.name1, self.name2, self.name3], [self.ress1, self.ress2, self.ress3], self.ressColumns)
                return filePath
        except IOError:
            return None
        except:
            if path.isfile(filePath) and not isOverwriting:
                remove(filePath)
            return None
        
    ##Funkcja odczytująca plik z 3 dopasowanymi sekwencjami i wyznaczająca dopasowanie korzystając z aktualnie ustawionej macierzy podobieństwa i wartości kary
    #@param filePath ścieżka do pliku
    #@return wartość dopasowania, gdy poprawny plik, None gdy niepoprawny
    def readAlignment(self, filePath):
        try:
            with open(filePath, 'r') as fp:
                nameList, seqList = FastaParser.read_clustal(fp)
                #jezeli jest za duzo lub za malo sekwencji
                if len(nameList) != 3:
                    return None

                #jezeli podano pusta sekwencje lub zawierajaca znaki spoza NUCLEOTIDE_LIST
                for seq in seqList:
                    if seq == "" or any(map((lambda x: x not in SequenceAlignment.NUCLEOTIDE_LIST), seq)):
                        return None
                seq1 = seqList[0]
                seq2 = seqList[1]
                seq3 = seqList[2]
                
                alignmentVal = 0
                matrix = copy.copy(self.matrix)
                for i in range(len(seq1)):
                    alignmentVal += self._getAlignmentVal(matrix, self._getSignVal(seq1[i]), self._getSignVal(seq2[i]), self._getSignVal(seq3[i]))
                
                return alignmentVal
        except:
            return None
            
    ##Funkcja uzupełniająca wartości w macierzy podobieństw dla przerw
    def _fillGapInMatrix(self):
        for i in range(SequenceAlignment.DIM):
            self.matrix[SequenceAlignment.DIM-1][i] = self.matrix[i][SequenceAlignment.DIM-1] = self.penaltyVal
        self.matrix[SequenceAlignment.DIM-1][SequenceAlignment.DIM-1] = 0
        
    ##Funkcja odczytująca wartość kary i numer wybranej macierzy
    #@return: tuplę (kara, numer macierzy)
    def getParams(self):
        return (self.penaltyVal, self.matrixNum)
    
    ##Funkcja ustawiająca nową wartość kary i numer wybranej macierzy
    #@param penaltyVal wartość kary za przerwę
    #@param matrixNum indeks wybranej macierzy
    #@param matrix macierz, jeżeli wybrano zdefiniowaną przez użytkownika   
    def setParams(self, penaltyVal, matrixNum, matrix = None):
        #zła wartość kary lub nieprawidłowy indeks macierzy podobieństwa
        if penaltyVal > SequenceAlignment.PENALTY_MAX or matrixNum < SequenceAlignment.MIN_MATRIX_NUM or matrixNum > SequenceAlignment.MAX_MATRIX_NUM:
            return
        
        #podano indeks macierzy definiowanej przez użytkownika, a nie podano macierzy
        if matrixNum == SequenceAlignment.USER and matrix == None:
            return

        self.penaltyVal = penaltyVal
        self.matrixNum = matrixNum
        
        if self.matrixNum == SequenceAlignment.IUB:
            self.matrix = copy.copy(SequenceAlignment.IUB_MATRIX)
        elif self.matrixNum == SequenceAlignment.CLUSTALW:
            self.matrix = copy.copy(SequenceAlignment.CLUSTALW_MATRIX)
        else:
            self.matrix = matrix
        
        self._fillGapInMatrix()
    
#         #print("Model params set:", self.penaltyVal, self.matrixNum, self.matrix)
        
    ##Funkcja odczytująodczyt indeksu danego nukleotydu w liscie
    #@param sign poszukiwany znak 
    def _getSignVal(self, sign):
        #zabezpieczone przez parser - niepoprawne znaczki sie nie pojawia
        return SequenceAlignment.NUCLEOTIDE_LIST.index(sign)
    
    ##Funkcja odczytująca wartość kary dla trzech elementów i podanej macierzy
    def _getAlignmentVal(self, matrix, val1, val2, val3):
        return matrix[val1][val2] + matrix[val1][val3] + matrix[val2][val3];
    
    def _maxInd(self, vec):
        ind = argmax(vec)
        return (vec[ind], ind);
    
    ##Funkcja inicjalizująca trójwymiarowe macierze przystosowania i ścieżek
    def _initializeAlignmentMatrices(self, matrix, alignmentTab, alignmentPath):
        alignmentTab[0][0][0] = 0;
        alignmentPath[0][0][0] = -1;
        p = self._getSignVal('-');

        #inicjalizacja macierzy dopasowania i macierzy kierunku zgodnie z porównywanymi sekwencjami
        for i in range(1, len(self.seq1)+1):
            alignmentTab[i][0][0] = alignmentTab[i-1][0][0] + self._getAlignmentVal(matrix, self._getSignVal(self.seq1[i-1]),p,p);
            alignmentPath[i][0][0] = 3;

        for i in range(1, len(self.seq2)+1):
            alignmentTab[0][i][0] = alignmentTab[0][i-1][0] + self._getAlignmentVal(matrix, p,self._getSignVal(self.seq2[i-1]),p);
            alignmentPath[0][i][0] = 5;

        for i in range(1, len(self.seq3)+1):
            alignmentTab[0][0][i] = alignmentTab[0][0][i-1] + self._getAlignmentVal(matrix, p,p,self._getSignVal(self.seq3[i-1]));
            alignmentPath[0][0][i] = 6;

    ##Funkcja wyznaczająca macierz trójwymiarową najlepszego dopasowania oraz trójwymiarową macierz ścieżek
    def _findBestAlignment(self, matrix, alignmentTab, alignmentPath):
        seq1 = self.seq1
        seq2 = self.seq2
        seq3 = self.seq3
        p = self._getSignVal('-');
        c = -inf;

        for k in range(0, len(self.seq3)+1):
#             #print( "model k:", k)
            self.emit(SIGNAL("progressValue(int, int)"), k, len(self.seq3)+1)
            for j in range(0, len(self.seq2)+1):
                for i in range(0, len(self.seq1)+1):
                    if self.checkIfContinue():
                        return
                    if (i+j+k == 0) or (i==0 and j==0) or (i==0 and k==0) or (j==0 and k==0):
                        continue;

                    val1 = val2 = val3 = p;
                    if i > 0: val1 = self._getSignVal(seq1[i-1]);
                    if j > 0: val2 = self._getSignVal(seq2[j-1]);
                    if k > 0: val3 = self._getSignVal(seq3[k-1]);
                    
                    #sprawdzenie, czy w danym kierunku nastąpiło przemieszczenie, czy będzie to przerwa
                    if i==0:
                        (maxv, ind) = self._maxInd([c,c,c,c,alignmentTab[i][j-1][k-1] + self._getAlignmentVal(matrix, p, val2, val3),
                                                    alignmentTab[i][j-1][k] + self._getAlignmentVal(matrix, p, val2, p),
                                                    alignmentTab[i][j][k-1] + self._getAlignmentVal(matrix, p, p, val3)]);
                    elif j==0:
                        (maxv, ind) = self._maxInd([c,c,alignmentTab[i-1][j][k-1] + self._getAlignmentVal(matrix, val1, p, val3),
                                                    alignmentTab[i-1][j][k] + self._getAlignmentVal(matrix, val1, p, p),
                                                    c,c,alignmentTab[i][j][k-1] + self._getAlignmentVal(matrix, p, p, val3)]);
                    elif k==0:
                        (maxv, ind) = self._maxInd([c,alignmentTab[i-1][j-1][k] + self._getAlignmentVal(matrix, val1, val2, p), c,
                                                    alignmentTab[i-1][j][k] + self._getAlignmentVal(matrix, val1, p, p),c,
                                                    alignmentTab[i][j-1][k] + self._getAlignmentVal(matrix, p, val2, p),c]);
                    else:
                        (maxv, ind) = self._maxInd([alignmentTab[i-1][j-1][k-1] + self._getAlignmentVal(matrix, val1, val2, val3),
                                                    alignmentTab[i-1][j-1][k] + self._getAlignmentVal(matrix, val1, val2, p),
                                                    alignmentTab[i-1][j][k-1] + self._getAlignmentVal(matrix, val1, p, val3),
                                                    alignmentTab[i-1][j][k] + self._getAlignmentVal(matrix, val1, p, p),
                                                    alignmentTab[i][j-1][k-1] + self._getAlignmentVal(matrix, p, val2, val3),
                                                    alignmentTab[i][j-1][k] + self._getAlignmentVal(matrix, p, val2, p),
                                                    alignmentTab[i][j][k-1] + self._getAlignmentVal(matrix, p, p, val3)]);
                    alignmentTab[i][j][k] = maxv;
                    alignmentPath[i][j][k] = ind;
#                     #print('maxv,ind:', maxv, ind)
#         for a in range(0, len(self.seq3)+1):
#             for b in range(0, len(self.seq2)+1):
#                 for c in range(0, len(self.seq1)+1):
#                     #print(self.alignmentTab[c][b][a], ' ', end='')
#                 #print('')
#         #print('\n\n')
#         
#         for a in range(0, len(self.seq3)+1):
#             for b in range(0, len(self.seq2)+1):
#                 for c in range(0, len(self.seq1)+1):
#                     #print(self.alignmentPath[c][b][a], ' ', end='')
#                 #print('\n')
#         #print('\n')
        
    ##Funkcja odtwarza sciezke z macierzy alignmentPath, gdzie zapisane byly kierunki przejsc pomiedzy komorkami
    def _reconstructPath(self, alignmentPath):
        seq1 = self.seq1
        seq2 = self.seq2
        seq3 = self.seq3
        ress1 = [];
        ress2 = [];
        ress3 = [];
        i = len(seq1)
        j = len(seq2)
        k = len(seq3)
        while 1:
            if self.checkIfContinue():
                return
            
            #podążanie po macierzy dopasowania zgodnie z informacjami zapisanymi w macierzy kierunków
            direction = int(alignmentPath[i][j][k]);
            if direction == -1:
                ress1 = ress1[::-1]
                ress2 = ress2[::-1]
                ress3 = ress3[::-1]
                return (ress1, ress2, ress3)
            elif direction == 0:
                ress1.append(seq1[i-1]); i=i-1;
                ress2.append(seq2[j-1]); j=j-1;
                ress3.append(seq3[k-1]); k=k-1;
            elif direction == 1:
                ress1.append(seq1[i-1]); i=i-1;
                ress2.append(seq2[j-1]); j=j-1;
                ress3.append('-');
            elif direction == 2:
                ress1.append(seq1[i-1]); i=i-1;
                ress2.append('-');
                ress3.append(seq3[k-1]); k=k-1;
            elif direction == 3:
                ress1.append(seq1[i-1]); i=i-1;
                ress2.append('-');
                ress3.append('-');
            elif direction == 4:
                ress1.append('-');
                ress2.append(seq2[j-1]); j=j-1;
                ress3.append(seq3[k-1]); k=k-1;
            elif direction == 5:
                ress1.append('-');
                ress2.append(seq2[j-1]); j=j-1;
                ress3.append('-');
            elif direction == 6:
                ress1.append('-');
                ress2.append('-');
                ress3.append(seq3[k-1]); k=k-1;
                

#     @profile
    ##Funkcja uruchamiająca procedurę wyznaczająca dopasowanie trzech wczytanych sekwencji.
    #@return tuplę (sek1, sek2, sek3) dopasowanych sekwencji, listę z '*' i ' ' tam, gdzie nukleotydy są wszystkie takie same, wartość dopasowania
    def alignSequences(self):
        self.ress1 = []
        self.ress2 = []
        self.ress3 = []
        self.ressColumns = []
        #usunięcie warunku stopu
        self.stopProcessingMutex.lock()
        self.stopProcessingVal = 0
        self.stopProcessingMutex.unlock()
        if self.seq1:
            matrix = copy.copy(self.matrix)
            shape = (len(self.seq1)+1, len(self.seq2)+1, len(self.seq3)+1)
            #inicjalizacja macierzy dopasowania oraz macierzy kierunku
            alignmentTab = ndarray(shape=shape)#empty(shape=shape)
            alignmentPath = ndarray(shape=shape)#empty(shape=shape)
            self._initializeAlignmentMatrices(matrix, alignmentTab, alignmentPath)
            self._findBestAlignment(matrix, alignmentTab, alignmentPath)
            if self.checkIfContinue():
                return
            (self.ress1, self.ress2, self.ress3) = self._reconstructPath(alignmentPath)
            if self.checkIfContinue():
                return
            self.alignmentValue = alignmentTab[len(self.seq1)][len(self.seq2)][len(self.seq3)]
            self.findExactOutput()
            return (self.ress1, self.ress2, self.ress3), self.ressColumns, self.alignmentValue
        
    ##Funkcja wyznaczająca listę z '*' i ' '
    def _findExactColumns(self, str1, str2, str3, res):
        mini = min(len(str3), min(len(str1), len(str2)))
        maxi = max(len(str3), max(len(str1), len(str2)))
        
        for i in range(mini):
            if str1[i] == str2[i] == str3[i]:
                res.append('*')
            else:
                res.append(' ')
                
        for i in range(mini, maxi):
            res.append(' ')
    
    ##Funkcja wyznaczająca listę '*' i ' ' dla sekwencji wejściowych
    def findExactInput(self):
        self._findExactColumns(self.seq1, self.seq2, self.seq3, self.seqColumns)
    
    ##Funkcja wyznaczająca listę '*' i ' ' dla sekwencji dopasowanych
    def findExactOutput(self):
        self._findExactColumns(self.ress1, self.ress2, self.ress3, self.ressColumns)
    
    ##Funkcja sprawdzająca, czy nie wybrano z GUI zakończenia obliczania dopasowania
    #@return True, jeżeli należy przerwać dopasowywanie, False jeżeli kontynuować
    def checkIfContinue(self):
        self.stopProcessingMutex.lock()
        stopProcessingVal = self.stopProcessingVal
        self.stopProcessingMutex.unlock()
        
        if stopProcessingVal == 1:
            return True
        else:
            return False
        
    ##Funkcja, która wywołana przez GUI powoduje kulturalne zatrzymanie wątka obliczającego dopasowanie
    def stopProcessing(self):
        self.stopProcessingMutex.lock()
        self.stopProcessingVal = 1
        self.stopProcessingMutex.unlock()
        
    ##Funkcja definiująca zachowanie wątka - uruchomienie procedury dopasowania sekwencji
    def run(self):
        alignSeqTime = timeit.timeit(stmt=self.alignSequences, number=1)
        print 'Sequences alignment lasted: ' + str(alignSeqTime)
#         self.alignSequences()
    